"""
-*- coding: utf-8 -*-

@author: Du Changping
@time: 2021/11/7 21:08
@file name: create_vox_file
@software：PyCharm

Do not smash your computer!

"""

import os
import time
import argparse
import itk
import SimpleITK as sitk
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm

import tomophantom
from tomophantom.TomoP2D import Objects2D
from tomophantom.TomoP3D import Objects3D


def path_setting(model_name_: str):
    """

    :param model_name_: model name to create the folder
    :return: current dir, the folder path to save the vox file
    """
    cur_dir_ = os.path.dirname(__file__)
    model_dir_path_ = os.path.join(cur_dir_, model_name_)
    if not os.path.exists(model_dir_path_):
        os.mkdir(model_dir_path_)
    return cur_dir_, model_dir_path_


def add_timestamp():
    time_stamp_ = time.strftime("%m%d-%H%M", time.localtime())
    return time_stamp_


def define_2d_objects(n_size_: int = 256):
    """
    use .dat file to create the object
    notice both methods in define_2d_objects and
    define_3d_objects can generate the wanted objects
    :param n_size_: generated vox matrix dimension
    :return: np.ndarray
    """
    dat_file_path_ = "synthetic2d.dat"
    dat_file_path_ = os.path.join(os.path.dirname(__file__), dat_file_path_)
    object_2d_ = tomophantom.TomoP2D.Model(1, n_size_,
                                           dat_file_path_)
    return object_2d_


def define_3d_objects(n_size_: int = 256, object_3d_compose_=None):
    """
    use funtion to create the object,
    :param n_size_: generated vox matrix dimension
    :param object_3d_compose_: compose can be feed into, by default is a simple ELLIPCYLINDER model
    :return: np.ndarray
    """
    if object_3d_compose_ is not None:
        pass
    else:
        p1_ = {
            "Obj": Objects3D.ELLIPCYLINDER,
            # "Obj": Objects3D.CUBOID,
            "C0": 1,
            "x0": 0.0,
            "y0": 0.0,
            "z0": 0.0,
            "a": 0.66,
            "b": 0.66,
            "c": 0.66,
            "phi1": 90
        }
        object_3d_compose_ = [p1_]
    object_3d_ = tomophantom.TomoP3D.Object(n_size_, object_3d_compose_)
    return object_3d_


def save_nii_file(save_matrix_: np.ndarray, save_path_: str):
    assert save_path_.split(".")[-1] == "nii", "here I save the file in .nii file"
    save_image_ = sitk.GetImageFromArray(save_matrix_)
    sitk.WriteImage(save_image_, save_path_)


def create_nii_file(model_name_: str, object_dim_: int = 3,
                    n_size_: int = 96, object_3d_compose_=None,
                    time_stamp_=""):
    print("*"*60)
    print("start to generate the .nii file")
    cur_dir_, model_save_path_ = path_setting(model_name_)
    if object_dim_ == 2:
        target_object_ = define_2d_objects(n_size_)
    elif object_dim_ == 3:
        target_object_ = define_3d_objects(n_size_, object_3d_compose_)
    else:
        raise RuntimeError("dimension must be 2 either or 3")

    nii_save_path_ = os.path.join(model_save_path_, model_name_ + "_" + time_stamp_ + ".nii")
    save_nii_file(target_object_, nii_save_path_)
    print(nii_save_path_)
    print("{}.nii generation done".format(model_name_))
    print("*"*60)
    print("\n")


def create_3d_vox_file(scale_: float, model_name_: str,
                       pixel_dist: float, time_stamp_=""):
    cur_dir_, model_save_dir_path_ = path_setting(model_name_)
    nii_file_path_ = os.path.join(model_save_dir_path_, model_name_ + "_" + time_stamp_ + ".nii")
    phantom_ = sitk.ReadImage(nii_file_path_)
    phantom_ = sitk.GetArrayFromImage(phantom_)
    dim_ = phantom_.shape
    assert len(dim_) > 2, "here we only support 3d object"

    nx_ = dim_[0]  # matrix dimension
    ny_ = dim_[1]
    nz_ = dim_[2]

    dx_ = pixel_dist * scale_
    dy_ = pixel_dist * scale_
    dz_ = pixel_dist * scale_

    phantom_ = phantom_.flatten(order="C")

    total_voxels_ = phantom_.size

    log_path_ = os.path.join(model_save_dir_path_, "{}_{}_{}.txt".format(model_name_, scale_, time_stamp_))
    log_f_ = open(log_path_, "w")
    log_f_.write("nii size is {}\n".format(dim_))
    log_f_.write("pixel distance is {}\n".format(pixel_dist))
    log_f_.write("target object length is {}\n".format(dim_[0]*pixel_dist))
    log_f_.close()

    save_name_ = "{}_{}_{}.vox".format(model_name_, scale_, time_stamp_)
    save_path_ = os.path.join(model_save_dir_path_, save_name_)

    print("*"*60)
    print("vox file save path is {}".format(save_path_))

    f_ = open(save_path_, "w")
    f_.write('[SECTION VOXELS phantomER]\n')
    f_.write(str(nx_) + ' ' + str(ny_) + ' ' + str(nz_) + ' No. OF VOXELS IN X,Y,Z\n')
    f_.write(str(dx_) + ' ' + str(dy_) + ' ' + str(dz_) + ' VOXEL SIZE (cm) ALONG X,Y,Z\n')
    f_.write('1 COLUMN NUMBER WHERE MATERIAL ID IS LOCATED\n')
    f_.write('2 COLUMN NUMBER WHERE THE MASS DENSITY IS LOCATED\n')
    f_.write('0 BLANK LINES AT END OF X,Y-CYCLES (1=YES,0=NO)\n')
    f_.write('[END OF VXH SECTION]\n')

    print("start to generate .vox file")
    print("phantom_ has following values:{}".format(np.unique(phantom_)))
    for i in tqdm(range(total_voxels_)):
        # todo user use this to define the material vox
        if phantom_[i] == 0:  # air
            f_.write('1 ' + str(0.00120479) + '\n')
        elif phantom_[i] == 1:  # water
            f_.write('2 ' + str(1.0) + '\n')
        else:
            # f_.write('1 ' + str(0.00120479) + '\n')
            raise RuntimeError("there is a value {} in phantom,"
                               "please check the code and notice "
                               "if there are mistakes value between"
                               "density and its no".format(phantom_[i]))
    f_.close()
    print("generation .vox file done")
    print("*"*60)
    print("\n")


if __name__ == "__main__":
    # parser = argparse.ArgumentParser(description="define the model name")
    # parser.add_argument("model_name", default="model1109", type=str)
    # parser.add_argument("n_size", default=64, type=int)
    # parser.add_argument("scale", default=0.85, type=float)
    # args = parser.parse_args()
    # print(args)

    # time_stamp = add_timestamp()  # optional
    time_stamp = ""  # optional
    # create_nii_file(model_name_=args.model_name, n_size_=args.n_size, time_stamp_=time_stamp)
    # create_3d_vox_file(scale_=args.scale,  model_name_=args.model_name,
    #                    pixel_dist=0.05)

    # create_nii_file(model_name_="model1112", n_size_=64, time_stamp_=time_stamp)
    create_3d_vox_file(scale_=1,  model_name_="medical_2",
                       pixel_dist=0.05, time_stamp_=time_stamp)

